package com.itjcloud.iot.gateway.filter.web;

import io.micrometer.tracing.Span;
import io.micrometer.tracing.Tracer;
import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

@RequiredArgsConstructor
@Component
@Slf4j
@Order(0)
public class WebResponseTracerFilter implements WebFilter {

    private final Tracer tracer;

    @NonNull
    @Override
    public Mono<Void> filter(@NonNull ServerWebExchange exchange, @NonNull WebFilterChain chain) {
        //设置响应 traceId
        return Mono.justOrEmpty(tracer)
                .map(Tracer::nextSpan)
                .map(Span::context)
                .doOnNext(traceContext -> exchange.getResponse().getHeaders()
                        .add("X-B3-TraceId", traceContext.traceId()))
                .then(chain.filter(exchange));
    }
}
